"""
构造FermiNet
"""

import numpy
import torch
from torch import nn




class ForwardLayer1(nn.Module):
    def __init__(self, nup, ndn):
        super().__init__()
        self.nup = nup
        self.ndn = ndn
        self.nel = self.nup + self.ndn
        #f^{α}_{i} = concatenate(h^α_i, g^up, g^dn, g^{αup}_i, g^{αdn}_i)
        #每一个f都是（5，4）且共享一个V，
        #V*f -> h^α_i, V应该是（1，5）
        self.V = nn.Parameter(torch.randn(4, 20), requires_grad=True)
        self.b = nn.Parameter(torch.zeros(1, 4))
        #为了保证反对称性，所有的粒子之间的操作都应该是等价？
        self.W = nn.Parameter(torch.randn(4, 4), requires_grad=True)
        self.c = nn.Parameter(torch.zeros(4))

    
    def forward(
            self,
            rupdn: torch.Tensor,
            prev_h_ai: torch.Tensor,
            prev_h_aibj: torch.Tensor
        ):
        '''前向'''
        self.h_ai = torch.zeros_like(prev_h_ai)
        self.h_aibj = torch.zeros_like(prev_h_aibj)
        #首先组装g
        #print(prev_h_ai[:self.nup, :])
        gup = torch.sum(prev_h_ai[:self.nup, :], dim=0, keepdim=False)/self.nup
        #print(prev_h_ai[self.nup:, :])
        gdn = torch.sum(prev_h_ai[self.nup:, :], dim=0, keepdim=False)/self.ndn
        for ai in range(self.nel):
            g_aiup = torch.sum(prev_h_aibj[ai, :self.nup, :], dim=0, keepdim=False)/self.nup
            g_aidn = torch.sum(prev_h_aibj[ai, self.nup:, :], dim=0, keepdim=False)/self.ndn
            f_ai = torch.cat([prev_h_ai[ai, :], gup, gdn, g_aiup, g_aidn])
            #print("f_ai", f_ai)
            #print(f_ai.shape)
            #
            #print(torch.matmul(self.V, f_ai).squeeze().shape)
            #print(self.h_ai[ai, :].shape)
            self.h_ai[ai, :] = torch.tanh(
                torch.matmul(self.V, f_ai) + self.b
            ).squeeze() + prev_h_ai[ai, :]
            #第二个指标中的粒子是可能交换位置的，所以要给他们相同的处理方式
            for bj in range(self.nel):
                self.h_aibj[ai, bj, :] = torch.tanh(
                    torch.matmul(self.W, prev_h_aibj[ai, bj, :]) + self.c
                ) + prev_h_aibj[ai, bj, :]
        return rupdn, self.h_ai, self.h_aibj



